import struct
import random
import numpy as np
import os
os.environ["MXNET_CPU_FORCE_F16"] = "1"

def fp16_to_hex(f):
    """将FP16数值转换为16进制字符串"""
    return f"{np.uint16(f):04X}"

def generate_fp16_case():
    """生成单个FP16测试用例（与原实现一致）"""
    num_type = random.choice(["normal", "normal", "denormal", "zero", "inf"])

    if num_type == "zero":
        sign = random.choice([0, 1])
        return (sign << 15) | 0x0000
    elif num_type == "inf":
        sign = random.choice([0, 1])
        return (sign << 15) | 0x7C00
    elif num_type == "denormal":
        sign = random.choice([0, 1])
        mantissa = random.randint(1, 0x3FF)
        return (sign << 15) | mantissa
    else:  # normal
        sign = random.choice([0, 1])
        exp = random.randint(1, 0x1E)  # 1-30
        mantissa = random.randint(0, 0x3FF)
        return (sign << 15) | (exp << 10) | mantissa

def fp16_add(a, b):
    """使用NumPy的FP16加法并处理特殊NaN情况"""
    np_a = np.uint16(a).view(np.float16)
    np_b = np.uint16(b).view(np.float16)
    sum_val = np_a + np_b

    # 统一NaN表示为静默NaN 0x7E00
    if np.isnan(sum_val):
        return 0x7E00
    return np.uint16(sum_val.view(np.uint16))

# 生成测试用例
test_cases = []
for _ in range(100):
    a = generate_fp16_case()
    b = generate_fp16_case()
    expected = fp16_add(a, b)

    test_cases.append({
        "a": fp16_to_hex(a),
        "b": fp16_to_hex(b),
        "expected": fp16_to_hex(expected)
    })

# 输出到文件
with open("fp16_add_cases.txt", "w") as f:
    f.write("// FP16加法测试用例（输入A，输入B，预期结果）\n")
    for case in test_cases:
        f.write(f"{case['a']} {case['b']} {case['expected']}\n")

print("生成100个测试用例到fp16_add_cases.txt")